package com.zjw.zy.heart.wbsocket;

import com.google.protobuf.InvalidProtocolBufferException;
import com.zjw.zy.heart.smi.JustWithIt;
import com.zjw.zy.heart.utils.LogUtils;
import com.zjw.zy.heart.wbsocket.socket.WsMessage;
import com.zjw.zy.pal.WsMessageWrap;
import lombok.NonNull;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ：zhong.jw
 * @date ：Created in 2022/12/16 11:15
 */
public final class WsStreamProcessor {

    private static final String TAG = "WsStreamProcessor";

    //最大传输count定位100w
    private static final int MAX_MESSAGE_TEXT_COUNT = 1_000_000;

    private static final Map<Long, List<WsMessageWrap.StreamItem>> streamMap = new ConcurrentHashMap<>();


    public static boolean sendByStream(@NonNull WsMessage message, @NonNull JustWithIt<byte[]> call) {
        if (message.getText() != null && message.getText().length() >= MAX_MESSAGE_TEXT_COUNT) {
            LogUtils.d(TAG, "streamSend() called with: message = [" + message.getId() + "]");
            WsMessageWrap.WsMessage.Builder wsBuilder = WsMessageWrap.WsMessage.newBuilder();

            wsBuilder.setInfo(message.getInfo());
            wsBuilder.setDate(message.getDate());
            wsBuilder.setFrom(message.getFrom());
            wsBuilder.setTo(message.getTo());
            wsBuilder.setId(message.getId());
            //stream-text
            String text = message.getText();

            //分流item
            WsMessageWrap.StreamItem.Builder itemBuilder = WsMessageWrap.StreamItem.newBuilder();
            //分流id和消息id设为相同
            itemBuilder.setId(message.getId());

            //分流个数
            int count = text.length() / MAX_MESSAGE_TEXT_COUNT;
            int index = 0;
            int end = 0;

            //分段总量
            itemBuilder.setCount(count + 1);
            while (index < count) {
                end = Math.min(text.length(), (1 + index) * MAX_MESSAGE_TEXT_COUNT);
                String item = text.substring(MAX_MESSAGE_TEXT_COUNT * index, end);
                wsBuilder.setText(item);

                itemBuilder.setItem(wsBuilder.build());
                //表示子段顺序
                itemBuilder.setSub(index);

                byte[] bytes = itemBuilder.build().toByteArray();
                call.onWith(bytes);
                index++;
            }

            //尾部数据
            if (end != text.length()) {
                wsBuilder.setText(text.substring(MAX_MESSAGE_TEXT_COUNT * index));
                itemBuilder.setItem(wsBuilder.build());
                itemBuilder.setSub(index);

                byte[] bytes = itemBuilder.build().toByteArray();
                call.onWith(bytes);
            }

            LogUtils.d(TAG, "streamSend() end " + message.getId());

            return true;
        }
        return false;
    }

    public static void acceptBytes(byte[] bytes, JustWithIt<WsMessage> call) {
        try {
            WsMessageWrap.StreamItem item = WsMessageWrap.StreamItem.parseFrom(bytes);
            processStreamItem(item, call);
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
    }


    private static void processStreamItem(WsMessageWrap.StreamItem item, JustWithIt<WsMessage> call) {
        List<WsMessageWrap.StreamItem> streamItems = streamMap.computeIfAbsent(item.getId(), aLong -> new ArrayList<>());
        streamItems.add(item);

        if (streamItems.size() == item.getCount()) {
            //分段完成，开始解析
            finishStream(streamItems, call);
        }
    }

    private static void finishStream(List<WsMessageWrap.StreamItem> streamItems, JustWithIt<WsMessage> call) {
        if (streamItems.isEmpty()) {
            return;
        }

        streamItems.sort(Comparator.comparingInt(WsMessageWrap.StreamItem::getSub));
        WsMessageWrap.StreamItem head = streamItems.get(0);
        WsMessage realMessage = new WsMessage();
        realMessage.setId(head.getId());
        realMessage.setInfo(head.getItem().getInfo());
        realMessage.setFrom(head.getItem().getFrom());
        realMessage.setTo(head.getItem().getTo());
        realMessage.setDate(head.getItem().getDate());

        StringBuilder text = new StringBuilder();
        for (WsMessageWrap.StreamItem streamItem : streamItems) {
            text.append(streamItem.getItem().getText());
        }

        realMessage.setText(text.toString());

        //移除缓存
        streamMap.remove(head.getId());

        call.onWith(realMessage);
    }

}
